import matplotlib.pyplot as plt
import numpy as np
import tensorflow as tf
from tensorflow.keras.preprocessing.image import ImageDataGenerator
from tensorflow.python.keras.callbacks import EarlyStopping

# 数据集处理：数据增强
train_datagen = ImageDataGenerator(rescale=1.0 / 255,
                                   rotation_range=40,
                                   width_shift_range=0.2,
                                   height_shift_range=0.2,
                                   shear_range=0.2,
                                   zoom_range=0.2,
                                   horizontal_flip=True,
                                   fill_mode='nearest')
# 标签生成，批次定义
train_generator = train_datagen.flow_from_directory('train',
                                                    batch_size=32,
                                                    classes=['with_mask', 'without_mask'],
                                                    target_size=(250, 250))
val_datagen = ImageDataGenerator(rescale=1.0 / 255)
val_generator = val_datagen.flow_from_directory('val',
                                                batch_size=8,
                                                classes=['with_mask', 'without_mask'],
                                                target_size=(250, 250))
model = tf.keras.models.Sequential([
    tf.keras.layers.Conv2D(32, (3, 3), activation='relu', input_shape=(250, 250, 3)),
    tf.keras.layers.MaxPooling2D(2, 2),
    tf.keras.layers.Conv2D(64, (3, 3), activation='relu'),
    tf.keras.layers.MaxPooling2D(2, 2),
    tf.keras.layers.Conv2D(128, (3, 3), activation='relu'),
    tf.keras.layers.MaxPooling2D(2, 2),
    tf.keras.layers.Flatten(),
    tf.keras.layers.Dropout(0.5),
    tf.keras.layers.Dense(50, activation='relu'),
    tf.keras.layers.Dense(2, activation='softmax')
])
model.compile(optimizer='adam', loss='binary_crossentropy', metrics=['acc'])

early_stopping = EarlyStopping(
    monitor='val_acc',
    verbose=1,
    patience=20,
    restore_best_weights=True
)

history = model.fit_generator(train_generator,
                              epochs=20,
                              validation_data=val_generator,
                              callbacks=[early_stopping])
model.save('cnn.h5')
plt.figure(figsize=(15, 6))
my_x_ticks = np.arange(0, 20, 1)
plt.xticks(my_x_ticks)
# loss
plt.plot(history.epoch, history.history.get('loss'), color='green', label='loss')
# acc
plt.plot(history.epoch, history.history.get('val_loss'), color='red', label='val_loss')
plt.legend()
plt.show()
plt.figure(figsize=(15, 6))
my_x_ticks=np.arange(0,20,1)
plt.plot(history.epoch,history.history.get('acc'),color='green',label = 'acc')
# acc
plt.plot(history.epoch,history.history.get('val_acc'),color='red',label = 'val_acc')
plt.legend()
plt.show()